# K-Means Clustering (medium)
import numpy as np

# 编写一个 Python 函数，实现用于聚类的 k-Means 算法，从指定的初始质心和一定数量的迭代开始。
# 该函数应采用
#       一个点列表（每个点表示为一个坐标元组）、
#       一个整数 k（表示要形成的聚类数）、
#       一个初始质心列表（每个质心表示为一个坐标元组）
#       和一个整数（表示要执行的最大迭代次数）。
# 该函数将迭代地将每个点分配给最近的质心，
# 并根据分配更新质心，直到质心不再发生显著变化，或者达到最大迭代次数。
# 该函数应返回聚类的最终质心列表。
# 四舍五入到最接近的第四位小数。


def euclidean_distance(a, b):
    return np.sqrt(((a - b) ** 2).sum(axis=1))


def k_means_clustering(
    points: list[tuple[float, float]],
    k: int,
    initial_centroids: list[tuple[float, float]],
    max_iterations: int,
) -> list[tuple[float, float]]:
    # Your code here
    points = np.array(points)
    centroids = np.array(initial_centroids)
    for _ in range(max_iterations):
        # Assign points to the nearest centroid
        distances = np.array(
            [euclidean_distance(points, centroid) for centroid in centroids]
        )
        assignments = np.argmin(distances, axis=0)

        new_centroids = np.array(
            [
                (
                    points[assignments == i].mean(axis=0)
                    if len(points[assignments == i]) > 0
                    else centroids[i]
                )
                for i in range(k)
            ]
        )

        # Check for convergence
        if np.all(centroids == new_centroids):
            break
        centroids = new_centroids
        centroids = np.round(centroids, 4)
    return centroids.tolist()
